"""
Unit tests for InfluxDBWriterService
"""

import pytest
import asyncio
from unittest.mock import Mock, AsyncMock
from decimal import Decimal
from datetime import datetime, timezone

from app.services.core.influxdb_writer_service import (
    InfluxDBWriterService,
    WriterMetrics,
    create_writer_service,
)
from app.adapters.core import Bar, Tick, DataFrequency, TickDirection, MarketType, Exchange


class MockInfluxDBAdapter:
    """Mock InfluxDB adapter for testing"""
    
    def __init__(self):
        self.write_batch_calls = []
        self.fail_next_write = False
    
    async def write_batch(self, batch):
        """Mock write_batch method"""
        self.write_batch_calls.append(batch)
        
        if self.fail_next_write:
            self.fail_next_write = False
            raise Exception("Mock write failure")
        
        await asyncio.sleep(0.01)  # Simulate I/O


@pytest.fixture
def mock_adapter():
    """Fixture providing a mock InfluxDB adapter"""
    return MockInfluxDBAdapter()


@pytest.fixture
async def writer_service(mock_adapter):
    """Fixture providing a writer service"""
    service = InfluxDBWriterService(
        adapter=mock_adapter,
        batch_size=10,  # Small batch size for testing
        flush_interval=0.5,  # Short interval for testing
        max_retries=2,
    )
    await service.start()
    yield service
    await service.stop()


@pytest.mark.asyncio
class TestInfluxDBWriterService:
    """Tests for InfluxDBWriterService"""
    
    async def test_write_bars(self, writer_service, mock_adapter):
        """Test writing bars"""
        bars = [
            Bar(
                symbol="000001.SZ",
                timestamp=datetime.now(timezone.utc),
                open=Decimal("10.00"),
                high=Decimal("10.50"),
                low=Decimal("9.90"),
                close=Decimal("10.20"),
                volume=Decimal("1000"),
                frequency=DataFrequency.DAILY,
            )
            for _ in range(5)
        ]
        
        await writer_service.write_bars(bars)
        
        # Buffer should contain 5 points
        assert len(writer_service.buffer) == 5
        
        # Flush
        await writer_service.flush()
        
        # Buffer should be empty
        assert len(writer_service.buffer) == 0
        
        # Adapter should have received 1 batch
        assert len(mock_adapter.write_batch_calls) == 1
        assert len(mock_adapter.write_batch_calls[0]) == 5
    
    async def test_write_ticks(self, writer_service, mock_adapter):
        """Test writing ticks"""
        ticks = [
            Tick(
                symbol="000001.SZ",
                timestamp=datetime.now(timezone.utc),
                price=Decimal("10.50"),
                volume=Decimal("1000"),
                direction=TickDirection.BUY,
            )
            for _ in range(3)
        ]
        
        await writer_service.write_ticks(ticks)
        
        # Buffer should contain 3 points
        assert len(writer_service.buffer) == 3
        
        await writer_service.flush()
        assert len(mock_adapter.write_batch_calls) == 1
    
    async def test_auto_batch(self, writer_service, mock_adapter):
        """Test automatic batching when batch_size is reached"""
        bars = [
            Bar(
                symbol=f"00000{i}.SZ",
                timestamp=datetime.now(timezone.utc),
                open=Decimal("10.00"),
                high=Decimal("10.50"),
                low=Decimal("9.90"),
                close=Decimal("10.20"),
                volume=Decimal("1000"),
                frequency=DataFrequency.DAILY,
            )
            for i in range(15)  # More than batch_size (10)
        ]
        
        await writer_service.write_bars(bars)
        
        # Wait for auto-flush
        await asyncio.sleep(0.2)
        
        # Should have triggered auto-flush
        assert len(mock_adapter.write_batch_calls) >= 1
    
    async def test_auto_flush(self, mock_adapter):
        """Test automatic flush based on interval"""
        # Create writer with short flush interval
        service = InfluxDBWriterService(
            adapter=mock_adapter,
            batch_size=1000,  # Large batch size
            flush_interval=0.3,  # 300ms interval
        )
        await service.start()
        
        try:
            # Add a few points
            bars = [
                Bar(
                    symbol="000001.SZ",
                    timestamp=datetime.now(timezone.utc),
                    open=Decimal("10.00"),
                    high=Decimal("10.50"),
                    low=Decimal("9.90"),
                    close=Decimal("10.20"),
                    volume=Decimal("1000"),
                    frequency=DataFrequency.DAILY,
                )
                for _ in range(5)
            ]
            
            await service.write_bars(bars)
            
            # Wait for auto-flush
            await asyncio.sleep(0.5)
            
            # Should have flushed automatically
            assert len(mock_adapter.write_batch_calls) >= 1
            assert len(service.buffer) == 0
            
        finally:
            await service.stop()
    
    async def test_metrics(self, writer_service, mock_adapter):
        """Test metrics tracking"""
        bars = [
            Bar(
                symbol="000001.SZ",
                timestamp=datetime.now(timezone.utc),
                open=Decimal("10.00"),
                high=Decimal("10.50"),
                low=Decimal("9.90"),
                close=Decimal("10.20"),
                volume=Decimal("1000"),
                frequency=DataFrequency.DAILY,
            )
            for _ in range(3)
        ]
        
        await writer_service.write_bars(bars)
        await writer_service.flush()
        
        metrics = writer_service.get_metrics()
        
        assert metrics.total_points == 3
        assert metrics.successful_writes >= 1
        assert metrics.batch_count >= 1
        assert metrics.last_flush_time is not None
    
    async def test_retry_logic(self, mock_adapter):
        """Test retry logic on write failure"""
        service = InfluxDBWriterService(
            adapter=mock_adapter,
            batch_size=10,
            flush_interval=10.0,
            max_retries=2,
        )
        await service.start()
        
        try:
            # Inject failure
            mock_adapter.fail_next_write = True
            
            bars = [
                Bar(
                    symbol="000001.SZ",
                    timestamp=datetime.now(timezone.utc),
                    open=Decimal("10.00"),
                    high=Decimal("10.50"),
                    low=Decimal("9.90"),
                    close=Decimal("10.20"),
                    volume=Decimal("1000"),
                    frequency=DataFrequency.DAILY,
                )
            ]
            
            await service.write_bars(bars)
            await service.flush()
            
            # Should have retried and succeeded
            metrics = service.get_metrics()
            assert metrics.retry_count > 0
            
        finally:
            await service.stop()
    
    async def test_buffer_limit(self, writer_service, mock_adapter):
        """Test buffer size limit enforcement"""
        # Try to add more than max_buffer_size (default 10000)
        # We use a smaller limit in the fixture
        writer_service.max_buffer_size = 20
        
        bars = [
            Bar(
                symbol=f"00000{i}.SZ",
                timestamp=datetime.now(timezone.utc),
                open=Decimal("10.00"),
                high=Decimal("10.50"),
                low=Decimal("9.90"),
                close=Decimal("10.20"),
                volume=Decimal("1000"),
                frequency=DataFrequency.DAILY,
            )
            for i in range(30)  # More than limit
        ]
        
        await writer_service.write_bars(bars)
        
        # Wait for flushes
        await asyncio.sleep(0.2)
        
        # Buffer should not exceed limit
        assert len(writer_service.buffer) <= writer_service.max_buffer_size


@pytest.mark.asyncio
async def test_create_writer_service(mock_adapter):
    """Test create_writer_service helper function"""
    service = await create_writer_service(
        adapter=mock_adapter,
        batch_size=100,
        flush_interval=1.0,
        auto_start=True,
    )
    
    try:
        assert service is not None
        assert service._running is True
        
    finally:
        await service.stop()


if __name__ == "__main__":
    pytest.main([__file__, "-v"])

